// BitTwister.h
// Override of Mersenne Twister that improves performance by caching.
// Just as Mersenne Twister, this is not thread safe and shouldn't be used
// in cryptography.

#ifndef BITTWISTER_H
#define BITTWISTER_H

#include "MersenneTwister.h"

namespace st
{
	class BitTwister: public MTRand
	{
	protected:
		uint32_t _cachedRand; // Cached random number
		uint32_t _bitsUsable; // Number of bits usable from _cachedRand

		uint32_t GetMask(uint32_t bits)
		{
			// Variation of: http://stackoverflow.com/questions/1392059/algorithm-to-generate-bit-mask
			// Also see: http://graphics.stanford.edu/~seander/bithacks.html
			switch(bits)
			{
			case 0:
				return 0;
			case 32:
				return 0xFFFFFFFF;
			default:
				return (1U << (bits)) - 1U;
			}
		}

	public:
		BitTwister(): MTRand(), _cachedRand(0), _bitsUsable(0) {}
		BitTwister(const BitTwister& r): MTRand(r), _cachedRand(r._cachedRand), _bitsUsable(r._bitsUsable) {}
		BitTwister(const uint32_t oneSeed): MTRand(oneSeed), _cachedRand(0), _bitsUsable(0) {}
		BitTwister(uint32_t* const bigSeed, uint32_t const seedLength = MTRand::N): MTRand(bigSeed, seedLength), _cachedRand(0), _bitsUsable(0) {}

		uint32_t bitRandomInt(uint32_t bitsNeeded)
		{
			uint32_t res = 0;
			while(bitsNeeded > 0)
			{
				if(_bitsUsable == 0)
				{ // Cache empty, refill
					_cachedRand = randInt();
					_bitsUsable = 32;
				}

				if(bitsNeeded > _bitsUsable)
				{ // Use all the cached bits
					res = (res << bitsNeeded) | _cachedRand;
					bitsNeeded -= _bitsUsable;
					_bitsUsable = 0;
				}
				else
				{ // Use some of the cached bits
					uint32_t mask = GetMask(bitsNeeded);
					res = (res << bitsNeeded) | (_cachedRand & mask);
					_bitsUsable -= bitsNeeded;
					_cachedRand >>= bitsNeeded;
					bitsNeeded = 0;
				}
			}
			return res;
		}

		static bool Test()
		{
			BitTwister b;
			b._bitsUsable = 32;
			b._cachedRand = 0xFFFFFFFF;

			if(b.GetMask(5) != 31)
				return false;
			if(b.GetMask(1) != 1)
				return false;
			if(b.bitRandomInt(1) != 1)
				return false;
			if(b._bitsUsable != 31)
				return false;
			if(b.bitRandomInt(32) < 0x7FFFFFFF)
				return false;
			return true;
		}
	};
}

#endif  // BITTWISTER_H

